// Copyright 2022 jeff.li. and/or its affiliates.

#include <tbir/pipeline/global_unique_index.h>
#include <tbir/runtime/container_private.h>
#include <tbir/runtime/file_util.h>
#include <tbir/runtime/native_object_maker.h>
#include <tbir/runtime/native_object_registry.h>

#include <tbir/runtime/future_wrap.h>

namespace tbir::runtime {

    RTValue Future::get() const {
        return body_();
    }

    void Future::set_body(std::function<RTValue()> body) {
        this->body_ = std::move(body);
    }

    TBIR_REGISTER_NATIVE_OBJECT(Future)
            .SetConstructor([](PyArgs args) -> std::shared_ptr<void> {
                MXCHECK_LT(args.size(), 2) << "[Lazy Construction] Expect 0 or 1 arguments but get "
                                           << args.size();
                if (args.size() == 0) {
                    return std::make_shared<Future>();
                } else {
                    UserDataRef op = args[0].As<UserDataRef>();
                    auto lazy = std::make_shared<Future>();
                    lazy->set_body([op]() { return op.call(); });
                    return lazy;
                }
            })
            .RegisterFunction("get",
                              [](void *self, PyArgs args) -> RTValue {
                                  MXCHECK_EQ(args.size(), 0)
                                      << "[Future][func: get] Expect 0 arguments but get " << args.size();
                                  return reinterpret_cast<Future *>(self)->get();
                              })
            .RegisterFunction("__call__", [](void *self, PyArgs args) -> RTValue {
                MXCHECK_EQ(args.size(), 0) << "[Future][func: get] Expect 0 arguments but get "
                                           << args.size();
                return reinterpret_cast<Future *>(self)->get();
            });

    UserDataRef Future::make_future_udref(std::function<RTValue()> body) {
        auto udref = make_native_userdata("Future", {});
        Future *future =
                static_cast<Future *>((static_cast<NativeObject *>(udref.ud_ptr())->opaque_ptr_).get());
        future->set_body(std::move(body));
        return udref;
    }

}  // namespace tbir::runtime
